import numpy as np
from scipy import stats

# 参考网站：https://www.jb51.net/article/131976.htm
#          https://www.jianshu.com/p/1ff0f2e5a3ce

#硬币投掷结果
observations = np.array([[1,0,0,0,1,1,0,1,0,1],
            [1,1,1,1,0,1,1,1,0,1],
            [1,0,1,1,1,1,1,0,1,1],
            [1,0,1,0,0,0,1,1,0,0],
            [0,1,1,1,0,1,1,1,0,1]])


def emSingle(prior, observations):
    """
      EM算法的单次迭代
      Arguments
      ------------
      priors:[theta_A, theta_B]
      observation:[m X n matrix]

      Returns
      ---------------
      new_priors:[new_theta_A, new_theta_B]
      :param priors:
      :param observations:
      :return:
    """
    counts = {'A': {'H': 0, 'T': 0}, 'B': {'H': 0, 'T': 0}}
    theta_A = prior[0]
    theta_B = prior[1]

    # E-step
    for observation in observations:
        len_observation = len(observation) # 每组抛硬币的次数
        num_heads = observation.sum() # 正面朝上的次数
        num_tails = len_observation - num_heads # 反面朝上的次数
        # 二项分布求解公式  (因为抛硬币符合二项分布)
        contribution_A = stats.binom.pmf(num_heads, len_observation, theta_A)
        contribution_B = stats.binom.pmf(num_tails, len_observation, theta_B)

        weight_A = contribution_A / (contribution_A + contribution_B)
        weight_B = contribution_B / (contribution_A + contribution_B)

        # 更新在当前参数下A，B硬币产生的正反面次数
        counts['A']['H'] += weight_A * num_heads
        counts['A']['T'] += weight_A * num_tails
        counts['B']['H'] += weight_B * num_heads
        counts['B']['T'] += weight_B * num_tails

    # M-step
    new_theta_A = counts['A']['H'] / (counts['A']['H'] + counts['A']['T'])
    new_theta_B = counts['B']['H'] / (counts['B']['H'] + counts['B']['T'])
    return [new_theta_A, new_theta_B]



def em(observations, prior, tol=1e-6, iteras=50):
    """
     EM算法
     ：param observations :观测数据
     ：param prior：模型初值
     ：param tol：迭代结束阈值
     ：param iterations：最大迭代次数
     ：return：局部最优的模型参数
     """
    itera = 0
    while itera < iteras:
        new_prior = emSingle(prior, observations)
        delta_change = np.abs(new_prior[0] - prior[0])
        if delta_change < tol:
            break
        else:
            prior = new_prior
            itera += 1
    return prior, itera


if __name__ == '__main__':
    prior, itera = em(observations, [0.6, 0.5])
    print(prior, itera)
